// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/dwrf/writer/writer_context.h>
#include <pollux/common/compression/compression.h>
#include <pollux/exec/memory_reclaimer.h>

namespace kumo::pollux::dwrf {
namespace {
constexpr uint32_t MIN_INDEX_STRIDE = 1000;
}

WriterContext::WriterContext(
    const std::shared_ptr<const Config>& config,
    std::shared_ptr<memory::MemoryPool> pool,
    const dwio::common::MetricsLogPtr& metricLogger,
    const tz::TimeZone* sessionTimezone,
    const bool adjustTimestampToTimezone,
    std::unique_ptr<encryption::EncryptionHandler> handler)
    : config_{config},
      pool_{std::move(pool)},
      dictionaryPool_{
          pool_->addLeafChild(fmt::format("{}.dictionary", pool_->name()))},
      outputStreamPool_{
          pool_->addLeafChild(fmt::format("{}.compression", pool_->name()))},
      generalPool_{
          pool_->addLeafChild(fmt::format("{}.general", pool_->name()))},
      indexEnabled_{getConfig(Config::CREATE_INDEX)},
      indexStride_{getConfig(Config::ROW_INDEX_STRIDE)},
      compression_{getConfig(Config::COMPRESSION)},
      compressionBlockSize_{getConfig(Config::COMPRESSION_BLOCK_SIZE)},
      shareFlatMapDictionaries_{getConfig(Config::MAP_FLAT_DICT_SHARE)},
      stripeSizeFlushThreshold_{getConfig(Config::STRIPE_SIZE)},
      dictionarySizeFlushThreshold_{getConfig(Config::MAX_DICTIONARY_SIZE)},
      linearStripeSizeHeuristics_{
          getConfig(Config::LINEAR_STRIPE_SIZE_HEURISTICS)},
      streamSizeAboveThresholdCheckEnabled_{
          getConfig(Config::STREAM_SIZE_ABOVE_THRESHOLD_CHECK_ENABLED)},
      rawDataSizePerBatch_{getConfig(Config::RAW_DATA_SIZE_PER_BATCH)},
      // Currently logging with no metadata. Might consider populating
      // metadata with dwio::common::request::AccessDescriptor upstream and
      // pass down the metric log.
      metricLogger_{metricLogger},
      sessionTimezone_{sessionTimezone},
      adjustTimestampToTimezone_{adjustTimestampToTimezone},
      handler_{std::move(handler)} {
  const bool forceLowMemoryMode{getConfig(Config::FORCE_LOW_MEMORY_MODE)};
  const bool disableLowMemoryMode{getConfig(Config::DISABLE_LOW_MEMORY_MODE)};
  POLLUX_CHECK(!(forceLowMemoryMode && disableLowMemoryMode));
  checkLowMemoryMode_ = !forceLowMemoryMode && !disableLowMemoryMode;
  if (forceLowMemoryMode) {
    setLowMemoryMode();
  }
  if (handler_ == nullptr) {
    handler_ = std::make_unique<encryption::EncryptionHandler>();
  }
  validateConfigs();
  VKLOG(2) << fmt::format(
      "Compression config: {}", common::compressionKindToString(compression_));
}

WriterContext::~WriterContext() {
  releaseMemoryReservation();
}

void WriterContext::validateConfigs() const {
  // the writer is implemented with strong assumption that index is enabled.
  // Things like dictionary encoding will fail if not. Before we clean that up,
  // always require index to be enabled.
  POLLUX_CHECK(indexEnabled(), "index is required");
  if (indexEnabled()) {
    POLLUX_CHECK_GE(indexStride_, MIN_INDEX_STRIDE);
    // Java works with signed integer and setting anything above the int32_max
    // will make the java reader fail.
    POLLUX_CHECK_LE(indexStride_, INT32_MAX);
  }
  POLLUX_CHECK_GE(
      compressionBlockSize_, getConfig(Config::COMPRESSION_BLOCK_SIZE_MIN));
  POLLUX_CHECK_GE(
      getConfig(Config::COMPRESSION_BLOCK_SIZE_EXTEND_RATIO),
      dwio::common::MIN_PAGE_GROW_RATIO);
}

void WriterContext::initBuffer() {
  POLLUX_CHECK_NULL(compressionBuffer_);
  if (compression_ != common::CompressionKind_NONE) {
    compressionBuffer_ = std::make_unique<dwio::common::DataBuffer<char>>(
        *generalPool_, compressionBlockSize_ + PAGE_HEADER_SIZE);
  }
}

memory::MemoryPool& WriterContext::getMemoryPool(
    const MemoryUsageCategory& category) {
  switch (category) {
    case MemoryUsageCategory::DICTIONARY:
      return *dictionaryPool_;
    case MemoryUsageCategory::OUTPUT_STREAM:
      return *outputStreamPool_;
    case MemoryUsageCategory::GENERAL:
      return *generalPool_;
    default:
      POLLUX_FAIL("Unreachable: {}", static_cast<int>(category));
  }
}

int64_t WriterContext::getMemoryUsage(
    const MemoryUsageCategory& category) const {
  switch (category) {
    case MemoryUsageCategory::DICTIONARY:
      return dictionaryPool_->usedBytes();
    case MemoryUsageCategory::OUTPUT_STREAM:
      return outputStreamPool_->usedBytes();
    case MemoryUsageCategory::GENERAL:
      return generalPool_->usedBytes();
    default:
      POLLUX_FAIL("Unreachable: {}", static_cast<int>(category));
  }
}

int64_t WriterContext::getTotalMemoryUsage() const {
  return generalPool_->usedBytes() + dictionaryPool_->usedBytes() +
      outputStreamPool_->usedBytes();
}

int64_t WriterContext::availableMemoryReservation() const {
  return dictionaryPool_->availableReservation() +
      outputStreamPool_->availableReservation() +
      generalPool_->availableReservation();
}

int64_t WriterContext::releasableMemoryReservation() const {
  return generalPool_->parent()->releasableReservation();
}

int64_t WriterContext::releaseMemoryReservation() {
  const auto* aggregatePool = dictionaryPool_->parent();
  const auto beforeTotalReservation = aggregatePool->reservedBytes();
  dictionaryPool_->release();
  outputStreamPool_->release();
  generalPool_->release();
  const auto releasedMemory =
      beforeTotalReservation - aggregatePool->reservedBytes();
  POLLUX_CHECK_GE(releasedMemory, 0);
  return releasedMemory;
}

void WriterContext::abort() {
  compressionBuffer_.reset();
  physicalSizeAggregators_.clear();
  streams_.clear();
  dictEncoders_.clear();
  decodedVectorPool_.clear();
  decodedVectorPool_.shrink_to_fit();
  selectivityVector_.reset();
  releaseMemoryReservation();
}
} // namespace kumo::pollux::dwrf
